use std::{rc::Rc, mem::size_of, os::raw::c_void, iter::Map, collections::HashMap};

use ash::vk;

use crate::{vulkan::{vk_struct::{VkBuffer, VkMappedBuffer, VkPipeline, VkWriteDescriptorSetGenerator}, VkContext, vk_utils::utils_info}, VK_CHECK_UNSAFE, VK_LOGP, system::scene::SceneSystem};

pub mod mesh;

pub struct Mesh {
    pub vertex_buffer   : VkBuffer,
    pub index_buffer    : VkBuffer,
    pub index_count     : u32
}

pub struct RenderEntity {
    pub id                          : usize,

    pub mesh                        : Rc<Mesh>,

    pub per_mesh_uniform_data       : VkMappedBuffer<PerMeshUniformData>,
    pub per_mesh_descriptor_sets    : Vec<vk::DescriptorSet>
}

pub struct RenderScene {
    pub render_entities     : Vec<RenderEntity>,
    pub cached_meshes       : Vec<Rc<Mesh>>,
    pub entity_object_map   : HashMap<usize, usize>,

    pub entity_counter      : usize
}

pub struct PerMeshUniformData {
    pub model_mat   : cgmath::Matrix4<f32>
}

pub struct Vertex<T> {
    pub position    : cgmath::Vector3<T>
}

impl RenderScene {
    pub fn new(vk_context: &VkContext) -> Self {
        let vertex_constant = vec![
            Vertex::<f32>{position: cgmath::Vector3::new(  0.5,  0.5, -0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new( -0.5,  0.5, -0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new(  0.5, -0.5, -0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new( -0.5, -0.5, -0.5)},
            Vertex::<f32>{position: cgmath::Vector3::new(  0.5,  0.5,  0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new( -0.5,  0.5,  0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new(  0.5, -0.5,  0.5)}, 
            Vertex::<f32>{position: cgmath::Vector3::new( -0.5, -0.5,  0.5)}
        ]; 

        let index_constant = vec![
            0, 1, 2,
            1, 2, 3,
            4, 5, 6, 
            5, 6, 7, 
            0, 1, 4,
            1, 4, 5, 
            2, 3, 6, 
            3, 6, 7,
            0, 2, 4, 
            2, 4, 6, 
            1, 3, 5, 
            3, 5, 7
        ];

        let vertex_buffer           = vk_context.create_device_only_buffer((size_of::<Vertex<f32>>() * vertex_constant.len()) as u64, vk::BufferUsageFlags::VERTEX_BUFFER, Some(vertex_constant.as_ptr() as *const c_void));
        let index_buffer            = vk_context.create_device_only_buffer((size_of::<Vertex<f32>>() * index_constant.len())  as u64, vk::BufferUsageFlags::INDEX_BUFFER , Some(index_constant.as_ptr()  as *const c_void));

        let cube = Rc::new(Mesh{vertex_buffer, index_buffer, index_count: index_constant.len() as u32});

        Self {
            render_entities     : vec![],
            cached_meshes       : vec![cube],
            entity_object_map   : HashMap::new(),

            entity_counter      : 0
        }

    }

    fn setup_per_mesh_descriptor_set(vk_context: &VkContext, descriptor_set_layouts: &Vec<vk::DescriptorSetLayout>, mesh_uniform_buffer: &VkMappedBuffer<PerMeshUniformData>) -> Vec<vk::DescriptorSet> {
        // 根据输入的Layout创建一定数量的Desceriptor Set
        let mut descriptor_set_ai = utils_info::get_descriptor_set_allocate_info(); 
        descriptor_set_ai.descriptor_pool           = vk_context.vk_descriptor_pool;
        descriptor_set_ai.descriptor_set_count      = 1;
        descriptor_set_ai.p_set_layouts             = descriptor_set_layouts.as_ptr();
        // 创建Set
        let descriptor_sets = match VK_CHECK_UNSAFE!(vk_context.vk_device.allocate_descriptor_sets(&descriptor_set_ai)) {
            Some(inner_value)   => inner_value,
            None                => VK_LOGP!("Failed to allocate_descriptor_sets")
        };
        // 处理对应需要写入的Binding位置
        let mut write_descriptor_generator = VkWriteDescriptorSetGenerator::new();
        let descriptor_buffer_info = vk::DescriptorBufferInfo {
            buffer          : mesh_uniform_buffer.buffer.handler,
            offset          : 0,
            range           : size_of::<PerMeshUniformData>() as u64
        };
        write_descriptor_generator.add_buffer_info(&mesh_uniform_buffer.buffer.handler, &descriptor_sets[0], 0, size_of::<PerMeshUniformData>() as u64, descriptor_buffer_info);

        let write_descriptor_sets = write_descriptor_generator.generate();
        // 向Descriptor Set中写入对应的Handler
        unsafe{ vk_context.vk_device.update_descriptor_sets(write_descriptor_sets, &[]) }

        descriptor_sets
    }

    pub fn add_render_entity(&mut self, scene_id: usize, vk_context: &VkContext, descriptor_set_layouts: &Vec<vk::DescriptorSetLayout>) {
        let id = self.entity_counter;
        self.entity_object_map.insert(id, scene_id);
        self.entity_counter += 1;

        let mesh = self.cached_meshes[0].clone();
        let per_mesh_uniform_data       = vk_context.create_uniform_buffer::<PerMeshUniformData>();
        let per_mesh_descriptor_sets    = Self::setup_per_mesh_descriptor_set(vk_context, descriptor_set_layouts, &per_mesh_uniform_data);

        self.render_entities.push(RenderEntity { id, mesh, per_mesh_uniform_data, per_mesh_descriptor_sets });
    }

    pub fn update(&self, scene_system: &SceneSystem) {
        for render_entity in &self.render_entities {
            if let Some(mapped_handler) = render_entity.per_mesh_uniform_data.mapped_handler {
                let cpu_data = PerMeshUniformData{ model_mat: scene_system.get_object_model_mat(self.entity_object_map[&render_entity.id]) };
                unsafe { mapped_handler.copy_from_nonoverlapping(&cpu_data as *const PerMeshUniformData, 1) };
            }
        }
    }
}